package com.zjl.media.core;

import com.zjl.media.config.properties.MediaProperties;
import com.zjl.media.core.ffmpeg.FfmpegCommandEnum;
import com.zjl.media.core.ffmpeg.FfmpegContext;
import com.zjl.media.core.ffmpeg.FfmpegTask;
import com.zjl.media.core.javacv.JavacvContext;
import com.zjl.media.core.javacv.JavacvTask;
import com.zjl.media.netty.dto.VideoParam;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.util.CollectionUtils;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * @author linfuxin Created on 2022-03-01 15:46:35
 */
@Slf4j
public class MediaManager {
    /**
     * 推流任务上下文
     */
    private final ConcurrentHashMap<String, MediaTransferContext> CONTEXT_CACHE = new ConcurrentHashMap<>();

    private final ExecutorService executorService;
    private final MediaProperties mediaProperties;

    private static final boolean IS_WINDOWS = "\\".equals(File.separator);

    public MediaManager(ExecutorService executorService, MediaProperties mediaProperties) {
        this.mediaProperties = mediaProperties;
        this.executorService = executorService;
    }

    public synchronized void push(VideoParam videoParam, ChannelHandlerContext client) {
        if (!addClientIfExistTask(videoParam.getRtsp(), client)) {
            TransferEnum transferEnum = TransferEnum.getInstance(videoParam.isHardDecode(), videoParam.isJavacv());
            if (transferEnum == null) {
                throw new RuntimeException("不支持的转码方式");
            }
            MediaTransferTask mediaTransferTask = transferEnum == TransferEnum.SOFT_JAVACV ?
                    buildJavacvTask(videoParam, client) : buildFfmpegTask(videoParam, client);
            executorService.execute(mediaTransferTask);
            CONTEXT_CACHE.put(videoParam.getRtsp(), mediaTransferTask.getContext());
            log.info("添加推流任务,type={},url={},channelId={},当前客户端数={}", transferEnum, videoParam.getRtsp(), client.channel().id().asShortText(), mediaTransferTask.getContext().getClients().size());
        }
    }

    public synchronized void forceClearAll() {
        List<ChannelHandlerContext> clientToClose = new ArrayList<>();
        for (MediaTransferContext mediaTransferContext : CONTEXT_CACHE.values()) {
            if (!CollectionUtils.isEmpty(mediaTransferContext.getClients())) {
                clientToClose.addAll(mediaTransferContext.getClients());
                mediaTransferContext.getClients().clear();
            }
        }

        for (ChannelHandlerContext client : clientToClose) {
            log.info("关闭无用客户端!channelId={}", client.channel().id().asShortText());
            client.close();
        }
        log.debug("强制关闭{}个客户端!", clientToClose.size());
    }

    @Scheduled(fixedDelay = 30000)
    public void evictContextClient() {
        if (log.isDebugEnabled()) {
            log.debug("直播推送客户端定时检测开始!");
        }

        try {
            List<ChannelHandlerContext> clientToClose = new ArrayList<>();
            for (MediaTransferContext mediaTransferContext : CONTEXT_CACHE.values()) {
                if (CollectionUtils.isEmpty(mediaTransferContext.getClients())) {
                    continue;
                }
                List<ChannelHandlerContext> clientToRemove = new ArrayList<>();
                for (int i = 0, size = mediaTransferContext.getClients().size(); i < size; i++) {
                    ChannelHandlerContext client = mediaTransferContext.getClients().get(i);
                    if (client == null) {
                        // 异常数据
                        mediaTransferContext.getClients().remove(i);
                        continue;
                    }
                    if (!client.channel().isWritable()) {
                        clientToClose.add(client);
                        clientToRemove.add(client);
                    }
                }
                if (!CollectionUtils.isEmpty(clientToRemove)) {
                    mediaTransferContext.getClients().removeAll(clientToRemove);
                }
            }

            for (ChannelHandlerContext client : clientToClose) {
                log.info("关闭无用客户端!channelId={}", client.channel().id().asShortText());
                client.close();
            }
            if (log.isDebugEnabled()) {
                log.debug("直播推送客户端检测关闭{}个客户端!", clientToClose.size());
            }
        } catch (RuntimeException e) {
            log.warn("清理客户端缓存出现异常!", e);
        }
    }

    private JavacvTask buildJavacvTask(VideoParam videoParam, ChannelHandlerContext client) {
        JavacvContext javacvContext = new JavacvContext();
        javacvContext.setRtsp(videoParam.getRtsp());
        javacvContext.setOutput(new ByteArrayOutputStream());
        javacvContext.setClients(new CopyOnWriteArrayList<>(Collections.singletonList(client)));
        javacvContext.setHeaderLatch(new CountDownLatch(1));

        return new JavacvTask(javacvContext, mediaProperties);
    }

    private FfmpegTask buildFfmpegTask(VideoParam videoParam, ChannelHandlerContext client) {

        FfmpegContext ffmpegContext = new FfmpegContext();
        ffmpegContext.setRtsp(videoParam.getRtsp());
        ffmpegContext.setSocketTimeout(mediaProperties.getSocketTimeout());
        ffmpegContext.setBufferSize(mediaProperties.getBufferSize());
        ffmpegContext.setFfmpegCommandEnum(FfmpegCommandEnum.getInstance(IS_WINDOWS, mediaProperties.getGpu().getType(), videoParam.isHardDecode()));
        ffmpegContext.setOutput(new ByteArrayOutputStream());
        ffmpegContext.setClients(new CopyOnWriteArrayList<>(Collections.singletonList(client)));
        ffmpegContext.setHeaderLatch(new CountDownLatch(1));

        return new FfmpegTask(ffmpegContext);
    }

    private boolean addClientIfExistTask(String rtsp, ChannelHandlerContext client) {
        String clientIdToPush = client.channel().id().asShortText();
        // 推流任务已经在执行,不需要启动推流任务;将客户端加入推送集合
        if (CONTEXT_CACHE.containsKey(rtsp)) {

            // 客户端为空另起一个任务
            MediaTransferContext mediaTransferContext = CONTEXT_CACHE.get(rtsp);
            if (CollectionUtils.isEmpty(mediaTransferContext.getClients())) {
                return false;
            }

            // 已经在推
            if (mediaTransferContext.getClients().stream().anyMatch(i -> clientIdToPush.equals(i.channel().id().asShortText()))) {
                return true;
            }

            // 拉流器已经启动,但是还没有拉取第一帧
            if (mediaTransferContext.getHeader() == null) {
                try {
                    //noinspection ResultOfMethodCallIgnored
                    mediaTransferContext.getHeaderLatch().await(mediaProperties.getWaitHeaderFrameTime(), TimeUnit.SECONDS);
                    if (mediaTransferContext.getHeader() == null) {
                        throw new RuntimeException("拉流器读取header异常!");
                    }
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }

            // 写入拉流器的第一帧并添加客户端直接开始推送
            client.writeAndFlush(new BinaryWebSocketFrame(Unpooled.copiedBuffer(mediaTransferContext.getHeader())));
            mediaTransferContext.getClients().add(client);
            log.info("添加推流客户端, rtsp={}, channelId={}, 当前客户端数={}", rtsp, client.channel().id().asShortText(), mediaTransferContext.getClients().size());
            return true;
        }
        return false;
    }
}
